# -*- coding: utf-8 -*-
"""
:mod:`nubio.io`
===============


The ``nubio.io`` module provides parsers and writers for common biological 
data  file formats.

"""

from __future__ import (division, with_statement, absolute_import, generators,
                         nested_scopes)
from nubox.char import CharVector, CharArray
from nubox.const import NOCHILD
from .fasta import (fasta_seq_parser, fasta_aln_parser,
                    fasta_seq_writer, fasta_aln_writer)
from .pdb import (pdb_parser, pdb_writer)
from .stockholm import (stockholm_aln_parser, stockholm_aln_writer)
from itertools import chain
from copy import deepcopy

def _get_extension(fh):
    """
    (internal) return the extension from a file ``handle`` name.
    
    """
    name = fh.name
    return name.split('.')[-1]

def meta_mapper(vector, src, dst, map):
    """
    Maps keys between source and destination dictionaries. Useful for explicit
    migration of data from one dataformat to another.
    
    Arguments:
    
      - vector(``nubox.char.CharArray``) a subclass of ``nubox.char.CharArray``
      - src(``str``) key to source metadata dictionary
      - dst(``str``) key to destination metadata dictionary
      - map(sequence) sequence of key mappings
      
    """
    for kwargs in vector.childrenkwargs:
        src_dict = kwargs['meta'][src]
        dst_dict = kwargs['meta'][dst] = kwargs['meta'].get(dst, {})
        for src_key, dst_key in map:
            dst_dict[dst_key] = src_dict[src_key]

def seq_parser(fh, vector, **kwargs):
    """
    Generator over sequences in a file-like object. Yields instances of the
    given "vector".
    
    Arguments:
    
      - fh(``file``) A file-like object open for reading. 
      - vector(``nubox.char.CharVector``) a subclass of 
        ``nubox.char.CharVector``
    
    Additional keyworded arguments are deep-copied and passed to the "vector"
    constructor.
    
    """
    extension = _get_extension(fh)
    if extension == 'fasta':
        parser = fasta_seq_parser
    else:
        raise ValueError("file format unknown")
    for seq, meta in parser(fh):
        seq_kwargs = deepcopy(kwargs)
        try:
            seq_kwargs['meta'].update(meta)
        except KeyError:
            seq_kwargs['meta'] = meta
        yield vector(seq, **seq_kwargs)

def aln_parser(fh, vector, **kwargs):
    """
    Returns an alignment of sequences in a file-like object as an instance
    of the given "vector". No grand-children corresponding to symbols in the
    alignment are constructed.
    
    Arguments:
    
      - fh(``file``) A file-like object open for reading. 
      - vector(``nubox.char.CharArray``) a subclass of 
        ``nubox.char.CharArray``
    
    Additional keyworded arguments are updated with the inferred "shape", 
    and "childrenkwargs".
    
    """
    extension = _get_extension(fh)
    if extension == 'fasta':
        parser = fasta_aln_parser
    elif extension == 'stockholm':
        parser = stockholm_aln_parser
    else:
        raise ValueError("file format unknown")
    childs, child_metas, parent_meta = parser(fh)
    shape = (len(childs), len(childs[0]), len(vector.clsdummy))
    childrenkwargs = [{'meta':meta, 'child':NOCHILD} for meta in child_metas]
    kwargs['shape'] = shape
    kwargs['childrenkwargs'] = childrenkwargs
    parent_meta.update(deepcopy(kwargs.get('meta', ())))
    kwargs['meta'] = parent_meta
    v = vector("".join(childs), **kwargs)
    return v

def model_parser(fh, vector, **kwargs):
    """
    Parses a file or file-like ``object`` and yield macromolecular models 
    in a "vector" instance. Multiple models from a single file are supported.
    
    Arguments:
    
      - fh(``file``) file-like ``object`` open for reading. 
      - vector(``nubox.double.DoubleArray``) a subclass of 
        ``nubox.double.DoubleArray``
    
    Additional keyworded arguments are passed to the "vector" constructor.
    
    """
    extension = _get_extension(fh)
    if extension == 'pdb':
        parser = pdb_parser
    else:
        raise ValueError("file format unknown")
    for model_atoms, model_atoms_meta, model_meta in parser(fh): # (model_atoms, model_atoms_meta, model_meta)
        shape = (len(model_atoms), len(vector.clsdummy))
        childrenkwargs = [{'meta':meta, 'child':NOCHILD} for meta in model_atoms_meta]
        kwargs['shape'] = shape
        kwargs['childrenkwargs'] = childrenkwargs
        model_meta.update(deepcopy(kwargs.get('meta', ())))
        kwargs['meta'] = model_meta
        data = tuple(chain(*model_atoms))
        yield vector(data, **kwargs)
